import torch
import numpy as np


#############################################################
# 张量序列化
shape = (1, 16, 224, 224)           # 数据维度
dtype = 'float32'                     # 数据类型
t = eval(f'torch.randn(shape).to\
        (dtype=torch.{dtype})')     # 生成一个张量
nt = np.array(t.to('cpu'))          # 转化为numpy数组

with open('nt.bin', 'wb') as fp:    # 将nt转为二进制数据并写入文件
    fp.write(nt.tobytes())

with open('nt.bin', 'rb') as fp:    # 从二进制数据中还原
    buffer = fp.read()
    nt = np.frombuffer(buffer, np.dtype(dtype)).reshape(shape)
    t_ = torch.from_numpy(nt.copy()) # 使用copy防止警告


#############################################################
# 压缩和解压
import gzip

buf = open('nt.bin', 'rb').read()
zbuf = gzip.compress(buf)
xbuf = gzip.decompress(zbuf)
print(f"rate of compression: %{100 * len(zbuf) / len(xbuf) :.2f}")


#############################################################
# header
import json

header = {}
header['shape'] = shape
header['dtype'] = dtype
json.dump(header, open('header.json', 'w'))

# 还原
header_ = json.load(open('header.json'))
shape_, dtype_ = header_["shape"], header_["dtype"]
with open('nt.bin', 'rb') as fp:    # 从二进制数据中还原
    buffer = fp.read()
    nt = np.frombuffer(buffer, np.dtype(dtype_)).reshape(shape_)
    t_ = torch.from_numpy(nt.copy()) # 使用copy防止警告
